#!/usr/bin/env python

from itertools import izip
import optparse

import numpy

from physics.vector import Vector4
from physics.vector.array import Vector4Array, boosted

from charmicat import BadCommandLineError

def main():
    '''
    Usage: %prog [opts] args
    
    Description
    '''

    op = optparse.OptionParser(usage=main.__doc__)
    
    opts, args = op.parse_args()

    size = 10000

    px = numpy.random.uniform(1, 100, size=size)
    py = numpy.random.uniform(1, 100, size=size)
    pz = numpy.random.uniform(1, 100, size=size)

    m  = numpy.random.uniform(0.1, 1.0, size=size)
    
    e = numpy.sqrt(m**2 + px**2 + py**2 + pz**2)

    v4s = [Vector4(*tup) for tup in izip(px, py, pz, e)]
    bvs = [v4.BoostVector() for v4 in v4s]

    px2 = numpy.random.uniform(1, 100, size=size)
    py2 = numpy.random.uniform(1, 100, size=size)
    pz2 = numpy.random.uniform(1, 100, size=size)

    m2  = numpy.random.uniform(0.1, 1.0, size=size)
    
    e2 = numpy.sqrt(m2**2 + px**2 + py**2 + pz**2)

    v42s = [Vector4(*tup) for tup in izip(px2, py2, pz2, e2)]
    
    for v4, bv in izip(v42s, bvs):
        v4.Boost(bv)

    pxf = [v4.X() for v4 in v42s]
    pyf = [v4.Y() for v4 in v42s]
    pzf = [v4.Z() for v4 in v42s]
    ef = [v4.T() for v4 in v42s]
    
    v4a1 = Vector4Array(px, py, pz, e)
    bv = v4a1.boost_vector()
    v4a2 = Vector4Array(px2, py2, pz2, e2)

    v4a2.boost(bv)

    print "boost different vectors"
    if numpy.all(v4a2.x == pxf):
        print "X passes"

    if numpy.all(v4a2.y == pyf):
        print "Y passes"

    if numpy.all(v4a2.z == pzf):
        print "Z passes"

    if numpy.all(v4a2.t == ef):
        print "E passes"
    
    bv = v4s[0].BoostVector()

    v42s = [Vector4(*tup) for tup in izip(px2, py2, pz2, e2)]
    for v4 in v42s:
        v4.Boost(bv)

    pxf = [v4.X() for v4 in v42s]
    pyf = [v4.Y() for v4 in v42s]
    pzf = [v4.Z() for v4 in v42s]
    ef = [v4.T() for v4 in v42s]
    
    v4a1 = Vector4Array(px, py, pz, e)
    v4a2 = Vector4Array(px2, py2, pz2, e2)
    v4a2.boost(bv.X(), bv.Y(), bv.Z())

    print "boost one vector"
    if numpy.all(v4a2.x == pxf):
        print "X passes"

    if numpy.all(v4a2.y == pyf):
        print "Y passes"

    if numpy.all(v4a2.z == pzf):
        print "Z passes"

    if numpy.all(v4a2.t == ef):
        print "E passes"

    v4a2 = Vector4Array(px2, py2, pz2, e2)
    v4a2 = boosted(v4a2, bv)

    print "boosted"
    if numpy.all(v4a2.x == pxf):
        print "X passes"

    if numpy.all(v4a2.y == pyf):
        print "Y passes"

    if numpy.all(v4a2.z == pzf):
        print "Z passes"

    if numpy.all(v4a2.t == ef):
        print "E passes"


    return 0


if __name__ == '__main__':
    try:
        import sys
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
            
